import datetime
from typing import List

from sqlalchemy.orm import Session
from sqlalchemy import Column, DateTime, Date, JSON, String, BIGINT, TEXT, Index, Integer, and_, func

from common.database import Base
from common.database import SessionMixin


class YarnJobInfo(Base, SessionMixin):
    __tablename__ = 'bigdata_yarn_job_info'
    application_id = Column(String(60), nullable=True, comment="应用程序ID")
    resource_pool = Column(String(40), nullable=True, comment="资源池")
    user = Column(String(60), nullable=True, comment="用户")
    hive_query_id = Column(String(60), nullable=True, comment="HIVE查询ID")
    hive_query_string = Column(TEXT, nullable=True, comment="HIVE查询字符串")
    cpu_milliseconds = Column(BIGINT, nullable=True, comment="CPU时间")
    # mb_millis = Column(String(128), nullable=False, comment="总内存分配")
    # vcores_millis = Column(String(128), nullable=False, comment="CPU分配")
    # maps_total = Column(Integer, nullable=False, comment="Map任务数量")
    # reduces_total = Column(Integer, nullable=False, comment="Reduces任务数量")
    state = Column(String(20), nullable=True, comment="应用程序状态")
    start_time = Column(DateTime, nullable=True, comment="开始时间")
    end_time = Column(DateTime, nullable=True, comment="结束时间")
    job_response = Column(JSON, nullable=True, comment="任务结果")
    create_date = Column(Date, default=datetime.date.today, nullable=True, comment="创建日期")
    create_time = Column(DateTime, default=datetime.datetime.now, nullable=False, comment="创建时间")

    __table_args__ = (
        Index('idx_bigdata_yarn_job_info_uuid', 'uuid', unique=True),
        Index('idx_bigdata_yarn_job_info_application_id', 'application_id'),
        Index('idx_bigdata_yarn_job_info_create_date', 'create_date'),
        Index('idx_bigdata_yarn_job_info_create_time', 'create_time'),
        {
            'mysql_engine': 'innodb',
            'mysql_charset': 'utf8mb4',
            'mysql_collate': 'utf8mb4_unicode_ci',
            'extend_existing': True,
            'mysql_row_format': 'dynamic',
        }
    )

    def __repr__(self):
        return "<application_id='{} ,user='{} ,resource_pool='{}'>".format(self.application_id, self.user, self.resource_pool)

    @classmethod
    def get_duration_top_by_date(cls, date: str, db: Session, top: int = 10) -> List:
        return db.query(cls).with_entities(cls.application_id, cls.resource_pool, cls.user, cls.hive_query_id,
                                           (func.time_to_sec(cls.end_time) - func.time_to_sec(cls.start_time)).label("duration"),
                                           cls.cpu_milliseconds, cls.start_time, cls.end_time).filter(
            and_(cls.create_date == date, cls.cpu_milliseconds.isnot(None), cls.start_time.isnot(None), cls.end_time.isnot(None))) \
            .group_by(cls.application_id, cls.resource_pool, cls.user, cls.hive_query_id, cls.cpu_milliseconds,
                      cls.start_time, cls.end_time) \
            .order_by((func.time_to_sec(cls.end_time) - func.time_to_sec(cls.start_time)).label("duration").desc()).limit(top).to_dict()

    @classmethod
    def get_cpu_milliseconds_top_by_date(cls, date: str, db: Session, top: int = 10) -> List:
        return db.query(cls).with_entities(cls.application_id, cls.resource_pool, cls.user, cls.hive_query_id,
                                           (func.time_to_sec(cls.end_time) - func.time_to_sec(cls.start_time)).label("duration"),
                                           cls.cpu_milliseconds, cls.start_time, cls.end_time).filter(
            and_(cls.create_date == date, cls.cpu_milliseconds.isnot(None), cls.start_time.isnot(None), cls.end_time.isnot(None))) \
            .group_by(cls.application_id, cls.resource_pool, cls.user, cls.hive_query_id, cls.cpu_milliseconds,
                      cls.start_time, cls.end_time) \
            .order_by(cls.cpu_milliseconds.desc()).limit(top).to_dict()


if __name__ == '__main__':
    from pprint import pprint
    from common.database import SessionLocal

    session = SessionLocal()
    # kudu_table_info = impalaQueryInfo()
    # print(kudu_table_info.get_latest_time_range(session))
    pprint(YarnJobInfo.get_duration_top_by_date("2023-03-02", session))
